# 指定文件编码为UTF-8
# coding: utf-8


"""
代码执行流程说明：
加载并预处理MNIST数据集
定义训练函数，同时训练带BN和不带BN的两个相同结构的网络
测试16种不同的权重初始化尺度
对每种权重初始化：
训练两个网络(带BN和不带BN)
记录训练过程中的准确率
绘制准确率变化曲线
最终显示所有实验结果的可视化对比
- 加载并预处理MNIST数据集
- 定义训练函数，同时训练带BN和不带BN的两个相同结构的网络
- 测试16种不同的权重初始化尺度
- 对每种权重初始化：
  - 训练两个网络(带BN和不带BN)
  - 记录训练过程中的准确率
  - 绘制准确率变化曲线
- 最终显示所有实验结果的可视化对比

实验结果分析：
批量归一化使网络对权重初始化更鲁棒
在较大或较小的权重初始化下，BN网络都能保持较好的训练效果
不带BN的网络在极端权重初始化下可能无法正常训练
BN网络通常收敛更快，最终准确率更高
- 批量归一化使网络对权重初始化更鲁棒
- 在较大或较小的权重初始化下，BN网络都能保持较好的训练效果
- 不带BN的网络在极端权重初始化下可能无法正常训练
- BN网络通常收敛更快，最终准确率更高

注意事项：
批量归一化的效果在深层网络中更明显
学习率和批量大小会影响BN的效果
实验中使用少量数据是为了快速展示BN的优势
实际应用中BN通常配合其他优化器(如Adam)使用效果更好
- 批量归一化的效果在深层网络中更明显
- 学习率和批量大小会影响BN的效果
- 实验中使用少量数据是为了快速展示BN的优势
- 实际应用中BN通常配合其他优化器(如Adam)使用效果更好
"""


"""
batch_norm_test.py 功能说明：
1. 对比测试批量归一化(Batch Normalization)对神经网络训练的影响
2. 使用MNIST数据集，比较带BN和不带BN的5层神经网络训练效果
3. 测试不同权重初始化尺度(weight_init_std)下的训练表现
4. 可视化16种不同权重初始化情况下的训练准确率对比

关键点：
- 批量归一化可以缓解梯度消失/爆炸问题
- 使网络对权重初始化更鲁棒
- 加速模型收敛
- 提高训练稳定性
- 使用SGD优化器进行训练
"""

# 导入系统模块
import sys, os
# 添加父目录到系统路径
sys.path.append(os.pardir)
# 导入NumPy数值计算库
import numpy as np
# 导入matplotlib绘图库
import matplotlib.pyplot as plt
# 从dataset.mnist导入MNIST数据加载函数
from dataset.mnist import load_mnist
# 从common.multi_layer_net_extend导入扩展的多层网络类
from common.multi_layer_net_extend import MultiLayerNetExtend
# 从common.optimizer导入优化器类
from common.optimizer import SGD, Adam

# 加载MNIST数据集(自动归一化)
(x_train, t_train), (x_test, t_test) = load_mnist(normalize=True)

# 减少训练数据量到1000个样本(加速实验)
x_train = x_train[:1000]
t_train = t_train[:1000]

# 设置训练参数
max_epochs = 20      # 最大训练轮数
train_size = x_train.shape[0]  # 训练数据大小
batch_size = 100     # 批量大小
learning_rate = 0.01 # 学习率

# 定义训练函数
def __train(weight_init_std):
    """
    训练并比较带BN和不带BN的网络

    参数:
        weight_init_std: 权重初始化标准差

    返回:
        train_acc_list: 普通网络的训练准确率列表
        bn_train_acc_list: BN网络的训练准确率列表
    """
    # 初始化带BN的5层网络
    bn_network = MultiLayerNetExtend(input_size=784,
                                   hidden_size_list=[100, 100, 100, 100, 100],
                                   output_size=10,
                                   weight_init_std=weight_init_std,
                                   use_batchnorm=True)

    # 初始化普通5层网络(不带BN)
    network = MultiLayerNetExtend(input_size=784,
                                 hidden_size_list=[100, 100, 100, 100, 100],
                                 output_size=10,
                                 weight_init_std=weight_init_std)

    # 使用SGD优化器
    optimizer = SGD(lr=learning_rate)

    # 初始化准确率记录列表
    train_acc_list = []      # 普通网络准确率
    bn_train_acc_list = []   # BN网络准确率

    # 计算每轮的迭代次数
    iter_per_epoch = max(train_size / batch_size, 1)
    epoch_cnt = 0

    # 训练循环
    for i in range(1000000000):
        # 随机选择批量数据
        batch_mask = np.random.choice(train_size, batch_size)
        x_batch = x_train[batch_mask]
        t_batch = t_train[batch_mask]

        # 同时训练两个网络(带BN和不带BN)
        for _network in (bn_network, network):
            # 计算梯度
            grads = _network.gradient(x_batch, t_batch)
            # 更新参数
            optimizer.update(_network.params, grads)

        # 每个epoch结束时评估模型
        if i % iter_per_epoch == 0:
            # 计算普通网络准确率
            train_acc = network.accuracy(x_train, t_train)
            # 计算BN网络准确率
            bn_train_acc = bn_network.accuracy(x_train, t_train)

            # 记录准确率
            train_acc_list.append(train_acc)
            bn_train_acc_list.append(bn_train_acc)

            # 打印当前训练进度
            print("epoch:" + str(epoch_cnt) + " | " + str(train_acc) + " - " + str(bn_train_acc))

            epoch_cnt += 1
            # 达到最大epoch数时停止训练
            if epoch_cnt >= max_epochs:
                break

    return train_acc_list, bn_train_acc_list


# 主实验部分：测试不同权重初始化尺度
# 生成16个对数间隔的权重初始化尺度(从10^0到10^-4)
weight_scale_list = np.logspace(0, -4, num=16)
x = np.arange(max_epochs)

# 对每个权重尺度进行实验
for i, w in enumerate(weight_scale_list):
    print("============== " + str(i+1) + "/16" + " ==============")
    # 训练并获取准确率曲线
    train_acc_list, bn_train_acc_list = __train(w)

    # 绘制子图(4x4网格)
    plt.subplot(4,4,i+1)
    plt.title("W:" + str(w))

    # 只在最后一个子图添加图例
    if i == 15:
        plt.plot(x, bn_train_acc_list, label='Batch Normalization', markevery=2)
        plt.plot(x, train_acc_list, linestyle = "--", label='Normal(without BatchNorm)', markevery=2)
    else:
        plt.plot(x, bn_train_acc_list, markevery=2)
        plt.plot(x, train_acc_list, linestyle="--", markevery=2)

    # 设置y轴范围
    plt.ylim(0, 1.0)
    # 只显示最左侧子图的y轴标签
    if i % 4:
        plt.yticks([])
    else:
        plt.ylabel("accuracy")
    # 只显示底部子图的x轴标签
    if i < 12:
        plt.xticks([])
    else:
        plt.xlabel("epochs")
    # 添加图例
    plt.legend(loc='lower right')

# 显示图形
plt.show()
